import numpy
import os
from shutil import copy
import json

result_dir1 = '/home/leslie/project/fashion-iq/CosMo/result@1'
result_dir5 = '/home/leslie/project/fashion-iq/CosMo/result@5'
result_dir10 = '/home/leslie/project/fashion-iq/CosMo/result@10'
result_dir50 = '/home/leslie/project/fashion-iq/CosMo/result@50'
result_dir_neg = '/home/leslie/project/fashion-iq/CosMo/result@neg'

img_root = '/home/leslie/project/fashion-iq/data/images'
caption_path = '/home/leslie/project/fashion-iq/data/captions/cap.dress.val.json'


def extract(positive_samples_info_k):

    postive = {}
    for i in positive_samples_info_k:
        ref_idx = i['ref_idx']
        ref_filename = all_ref_attributes[ref_idx]
        targ_idxs = i['targ_idxs']
        targ_filenames = []
        for j in targ_idxs:
            targ_filenames.append(all_test_attributes[j])
        postive[ref_filename] = targ_filenames
    return postive


def mkdir(path):
    try:
        os.mkdir(path)
    except Exception as e:
        print(e)


def visual(positive_1, dir):
    result_dir = dir
    a = json.load(open(caption_path))
    query_ref_dic = {}
    for i in a:
        target = i['target']
        candidate = i['candidate']
        query_ref_dic[candidate] = target

    mkdir(result_dir)

    keys = list(positive_1.keys())
    for i in keys:

        sub_dir = os.path.join(result_dir, i)
        mkdir(sub_dir)
        filenames = positive_1[i]
        init = 1
        ground_truth = query_ref_dic[i]

        src = os.path.join(img_root, i + '.jpg')
        des = os.path.join(sub_dir, '_Query_' + i + '.jpg')
        copy(src, des)

        for file in filenames:

            if ground_truth != file:
                src = os.path.join(img_root, file + '.jpg')
                des = os.path.join(sub_dir, str(init).zfill(2) + '_' + file + '.jpg')
            else:
                src = os.path.join(img_root, file + '.jpg')
                des = os.path.join(sub_dir, str(init).zfill(2) + '_GroundTruth_' + file + '.jpg')
            init = init + 1
            copy(src, des)


if __name__ == '__main__':

    all_info = numpy.load('all_.npy', allow_pickle=True)
    all_info = all_info.tolist()
    all_ref_attributes = all_info['candidate']
    all_query_attributes = all_info['target']
    all_test_attributes = all_info['all']
    positive_samples_info = all_info['positive']
    positive_samples_info_1 = positive_samples_info[0]
    positive_samples_info_5 = positive_samples_info[1]
    positive_samples_info_10 = positive_samples_info[2]
    positive_samples_info_50 = positive_samples_info[3]
    negative_samples_info = all_info['negtive']

    positive_1 = extract(positive_samples_info_1)
    positive_5 = extract(positive_samples_info_5)
    positive_10 = extract(positive_samples_info_10)
    positive_50 = extract(positive_samples_info_50)
    negative = extract(negative_samples_info)
    visual(positive_1, result_dir1)
    visual(positive_5, result_dir5)
    visual(positive_10, result_dir10)
    visual(positive_50, result_dir50)
    visual(negative, result_dir_neg)

    print(1)
'''
'ref_idx'->all_ref_attributes  , 'targ_idxs'->all_test_attributes
'''
